import streamlit as st
import requests

# 设置页面配置
st.set_page_config(layout="wide")

# 设置页面样式
st.markdown("""
    <style>
        .conversation-container {
            max-height: 500px; /* 设置最大高度 */
            overflow-y: auto;
            padding-right: 1rem;
            display: flex;
            flex-direction: column; /* 纵向排列 */
            gap: 10px; /* 增加消息之间的间距 */
        }
        [data-testid="stVerticalBlock"] {
            gap: 0rem;
        }
        .conversation-container::-webkit-scrollbar {
            width: 8px;
        }
        .conversation-container::-webkit-scrollbar-track {
            background: #f1f1f1;
        }
        .conversation-container::-webkit-scrollbar-thumb {
            background: #888;
            border-radius: 4px;
        }
        .conversation-container::-webkit-scrollbar-thumb:hover {
            background: #555;
        }
        .user-message {
            margin: 5px;
            padding: 10px;
            border-radius: 10px;
            background-color: #e3f2fd; /* 浅蓝色 */
            align-self: flex-start; /* 左对齐 */
            display: flex;
            align-items: center;
            max-width: 70%;
        }
        .assistant-message {
            margin: 5px;
            padding: 10px;
            border-radius: 10px;
            background-color: #f5f5f5; /* 浅灰色 */
            align-self: flex-start; /* 左对齐 */
            display: flex;
            align-items: center;
            max-width: 70%;
        }
        .avatar {
            width: 30px;
            height: 30px;
            border-radius: 50%;
            margin-right: 10px; /* 头像在左边 */
        }
    </style>
""", unsafe_allow_html=True)

# 设置页面标题
st.title("RAG 智能问答系统")

# 创建左右两列布局
left_col, right_col = st.columns([1, 2])

# 左侧列 - 文档上传区域
with left_col:
    st.header("知识库文档上传")
    st.markdown("上传文档到知识库")
    
    uploaded_file = st.file_uploader(
        "支持 PDF、TXT、DOCX 格式",
        type=['pdf', 'txt', 'docx'],
        help="单个文件大小限制为200MB"
    )
    
    if uploaded_file:
        files = {"file": uploaded_file}
        with st.spinner("正在上传文件..."):
            response = requests.post("http://localhost:9002/upload", files=files)
            if response.status_code == 200:
                st.success("文档上传成功！")
            else:
                st.error("文档上传失败！")

# 右侧列 - 对话区域
with right_col:
    st.header("智能问答")
    
    if "messages" not in st.session_state:
        st.session_state.messages = []

    with st.container():
        st.markdown('<div class="conversation-container">', unsafe_allow_html=True)
        
        for message in st.session_state.messages:
            if message["role"] == "user":
                st.markdown(
                    f'''
                    <div class="user-message">
                        <img src="https://cdn-icons-png.flaticon.com/512/3135/3135715.png" class="avatar"/>
                        {message["content"]}
                    </div>
                    ''',
                    unsafe_allow_html=True
                )
            else:
                st.markdown(
                    f'''
                    <div class="assistant-message">
                        <img src="https://cdn-icons-png.flaticon.com/512/4712/4712035.png" class="avatar"/>
                        {message["content"]}
                    </div>
                    ''',
                    unsafe_allow_html=True
                )
        
        st.markdown('</div>', unsafe_allow_html=True)

# 用户输入框
prompt = st.chat_input("请输入您的问题")

# 处理用户输入
if prompt:
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    with right_col:
        with st.container():
            st.markdown(
                f'''
                <div class="user-message">
                    <img src="https://cdn-icons-png.flaticon.com/512/3135/3135715.png" class="avatar"/>
                    {prompt}
                </div>
                ''',
                unsafe_allow_html=True
            )
        
            # 显示助手回答
            with st.spinner("思考中..."):
                response = requests.post(
                    "http://localhost:9002/chat",
                    json={"question": prompt},
                    stream=True
                )
                
                if response.status_code == 200:
                    full_response = ""
                    response_container = st.empty()  # 创建一个可更新的容器
                    for chunk in response.iter_lines():
                        if chunk:
                            part = chunk.decode('utf-8')
                            full_response += part
                            response_container.markdown(
                                f'''
                                <div class="assistant-message">
                                    <img src="https://cdn-icons-png.flaticon.com/512/4712/4712035.png" class="avatar"/>
                                    {full_response}
                                </div>
                                ''',
                                unsafe_allow_html=True
                            )
                    st.session_state.messages.append({"role": "assistant", "content": full_response})
                else:
                    st.error(f"获取回答失败！状态码：{response.status_code}")

    st.markdown("""
        <script>
            window.onload = function() {
                const container = document.querySelector('.conversation-container');
                container.scrollTop = container.scrollHeight;
            };
        </script>
    """, unsafe_allow_html=True)

# 添加页面底部说明
st.markdown("---")
st.markdown("💡 提示：您可以通过上传文档来扩充知识库，提高回答的准确性。")